from mcp.server.fastmcp import FastMCP, Context
import subprocess
import json
from typing import Optional

mcp = FastMCP("CRI-O MCP Server")

"""
面向 CRI-O 容器运行时软件的 MCP Server

所有 MCP tool 列表：
1. crio_service_control - 控制 crio 服务的启动、停止、重启和状态查询
   参数：action (str): 操作类型，可选值：start, stop, restart, status
   返回：统一 JSON 结构

2. crio_service_reload - 重新加载 crio 服务配置
   参数：无
   返回：统一 JSON 结构

3. crio_service_modify_log_level - 修改 crio 服务的日志级别配置
   参数：log_level (str): 新的日志级别
   返回：统一 JSON 结构

4. crio_wipe_service_status - 查询 crio-wipe 服务的状态
   参数：无
   返回：统一 JSON 结构
"""

@mcp.tool()
def crio_service_control(action: str) -> dict:
    """
    控制 crio 服务的启动、停止、重启和状态查询
    
    参数:
        action: 操作类型，可选值：start, stop, restart, status
    
    返回:
        JSON 结构: {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    valid_actions = ["start", "stop", "restart", "status"]
    if action not in valid_actions:
        return {
            "success": False,
            "command": f"systemctl {action} crio",
            "exit_code": 1,
            "stdout": "",
            "stderr": f"Invalid action: {action}. Must be one of {valid_actions}"
        }
    
    command = f"systemctl {action} crio"
    try:
        result = subprocess.run(command, shell=True, capture_output=True, text=True, timeout=30)
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out"
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def crio_service_reload() -> dict:
    """
    重新加载 crio 服务配置
    
    参数:
        无
    
    返回:
        JSON 结构: {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    commands = [
        "systemctl daemon-reload",
        "systemctl reload crio"
    ]
    
    all_stdout = ""
    all_stderr = ""
    exit_code = 0
    
    for cmd in commands:
        try:
            result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=30)
            all_stdout += f"Command: {cmd}\n{result.stdout}\n"
            all_stderr += f"Command: {cmd}\n{result.stderr}\n"
            if result.returncode != 0:
                exit_code = result.returncode
        except subprocess.TimeoutExpired:
            all_stderr += f"Command: {cmd}\nTimed out\n"
            exit_code = -1
        except Exception as e:
            all_stderr += f"Command: {cmd}\nError: {str(e)}\n"
            exit_code = -1
    
    return {
        "success": exit_code == 0,
        "command": " && ".join(commands),
        "exit_code": exit_code,
        "stdout": all_stdout,
        "stderr": all_stderr
    }

@mcp.tool()
def crio_service_modify_log_level(log_level: str) -> dict:
    """
    修改 crio 服务的日志级别配置
    
    参数:
        log_level: 新的日志级别
    
    返回:
        JSON 结构: {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    # 备份原始配置
    backup_cmd = "cp /usr/lib/systemd/system/crio.service /usr/lib/systemd/system/crio.service.backup"
    
    # 修改日志级别
    modify_cmd = f"sed -i 's|ExecStart=/usr/bin/crio|ExecStart=/usr/bin/crio --log {log_level}|g' /usr/lib/systemd/system/crio.service"
    
    # 重新加载配置
    reload_cmds = [
        "systemctl daemon-reload",
        "systemctl reload crio"
    ]
    
    commands = [backup_cmd, modify_cmd] + reload_cmds
    
    all_stdout = ""
    all_stderr = ""
    exit_code = 0
    
    for cmd in commands:
        try:
            result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=30)
            all_stdout += f"Command: {cmd}\n{result.stdout}\n"
            all_stderr += f"Command: {cmd}\n{result.stderr}\n"
            if result.returncode != 0:
                exit_code = result.returncode
        except subprocess.TimeoutExpired:
            all_stderr += f"Command: {cmd}\nTimed out\n"
            exit_code = -1
        except Exception as e:
            all_stderr += f"Command: {cmd}\nError: {str(e)}\n"
            exit_code = -1
    
    return {
        "success": exit_code == 0,
        "command": " && ".join(commands),
        "exit_code": exit_code,
        "stdout": all_stdout,
        "stderr": all_stderr
    }

@mcp.tool()
def crio_wipe_service_status() -> dict:
    """
    查询 crio-wipe 服务的状态
    
    参数:
        无
    
    返回:
        JSON 结构: {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    command = "systemctl status crio-wipe.service"
    try:
        result = subprocess.run(command, shell=True, capture_output=True, text=True, timeout=30)
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out"
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

if __name__ == "__main__":
    mcp.run()